/**
 * @Author: wzx
 * @Description:
 * @Version: 1.0.0
 * @Date: 2021/5/18 上午1:11
 *@Copyright: MIN-Group；国家重大科技基础设施——未来网络北大实验室；深圳市信息论与未来网络重点实验室
 */

package keymanager

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"vpn-management-server/logger"
	"vpn-management-server/security/crypto/sm2"
	"vpn-management-server/security/crypto/sm4"
)

type KeyManager struct {
	prk *sm2.Sm2PrivateKey
	pub *sm2.Sm2PublicKey
}

func (km *KeyManager) Init() {
	km.prk = &sm2.Sm2PrivateKey{}
	km.pub = &sm2.Sm2PublicKey{}
}

func New(pub string, pri string) *KeyManager {
	keyManager := &KeyManager{}
	keyManager.Init()
	err := keyManager.SetPriKey(pri)
	if err != nil {
		logger.Logger.Fatal("set privateKey fail")
	}
	err = keyManager.SetPubKey(pub)
	if err != nil {
		logger.Logger.Fatal("set publicKey fail,err:", err)
	}
	return keyManager
}

func (km *KeyManager) GenKeyPair() {
	km.prk, km.pub = sm2.GenKeyPair()
}

//得到公钥 类型格式-->字符串格式
func (km *KeyManager) GetPubkey() string {
	return string(km.pub.GetBytes())
}

//得到私钥 类型格式-->字符串格式
func (km *KeyManager) GetPriKey() string {
	return string(km.prk.GetBytes())
}

func (km *KeyManager) SetPubKey(data string) error {
	err := km.pub.SetBytes([]byte(data))
	if err != nil {
		return err
	}
	return nil
}

func (km *KeyManager) SetPriKey(data string) error {
	err := km.prk.SetBytes([]byte(data))
	if err != nil {
		return err
	}
	return nil
}

func (km *KeyManager) Encrypt(data string) ([]byte, error) {
	return km.pub.Encrypt(rand.Reader, []byte(data), nil)
}

func (km *KeyManager) Decrypt(data string) ([]byte, error) {
	return km.prk.Decrypt(rand.Reader, []byte(data), nil)
}

func (km *KeyManager) IsOnCurve() bool {
	return sm2.IsOnCurve(km.prk, km.pub)
}

//得到数字签名
func (km *KeyManager) Sign(text []byte) (string, error) {
	sig, err := km.prk.Sign(rand.Reader, text, nil)
	return base64.StdEncoding.EncodeToString(sig), err
}

//验证签名
func (km *KeyManager) Verify(text []byte, signature string, pubkey string) (bool, error) {
	t, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		fmt.Println("base64编码错误")
		return false, err
	}

	pub := sm2.Sm2PublicKey{}
	err = pub.SetBytes([]byte(pubkey))
	if err != nil {
		fmt.Println(err)
		return false, err
	}
	res, err := pub.Verify(text, t, nil)
	if err != nil {
		fmt.Println(err)
		return false, err
	}
	return res, nil
}

func SM4Encrypt(secretkey string, origData []byte) ([]byte, error) {
	return sm4.Sm4Ecb([]byte(secretkey), origData, sm4.ENC)
}
func SM4Decrypt(secretkey string, crypted []byte) ([]byte, error) {
	return sm4.Sm4Ecb([]byte(secretkey), crypted, sm4.DEC)
}
